import torch.nn as nn
import torch.nn.functional as Func


class LeNet(nn.Module):  # 继承父类nn.module
    def __init__(self):  # 初始化网络结构
        super(LeNet, self).__init__()           # 多继承用到的super函数
        self.conv1 = nn.Conv2d(3, 16, 5)        # 第一层--3通道到16通道，卷积核尺寸：5x️5
        # 这里给出卷积函数的一般用法：
        # torch.nn.Conv2d(in_channels, out_channels, kernel_size,
        #                 stride=1, padding=0, dilation=1, groups=1,
        #                 bias=True, padding_mode='zeros')
        self.pool1 = nn.MaxPool2d(2, 2)         # 最大池化--尺寸为2x2，跳步为2
        # 这里给出最大池化函数的一般用法：
        # MaxPool2d(kernel_size, stride)
        self.conv2 = nn.Conv2d(16, 32, 5)       # 第二层--16通道到32通道，卷积核尺寸：5x5
        self.pool2 = nn.MaxPool2d(2, 2)         # 最大池化--尺寸为2x2，跳步为2
        self.fc1 = nn.Linear(32 * 5 * 5, 120)   # 全连接层--32x5x5全连接至120
        # 这里给出全连接层到一般用法：
        # Linear(in_features, out_features, bias=True)
        self.fc2 = nn.Linear(120, 84)           # 全连接层--120全连接至84
        self.fc3 = nn.Linear(84, 10)            # 全连接层--84全连接至10

    def forward(self, x):  # 正向传播函数
        x = Func.relu(self.conv1(x))    # input 3x32x32  output 16x28x28 activateFunction:relu
        x = self.pool1(x)               # input 16x28x28 output 16x14x14
        x = Func.relu(self.conv2(x))    # input 16x14x14 output 32x10x10
        x = self.pool2(x)               # input 32x10x10 output 32x5x5
        x = x.view(-1, 32 * 5 * 5)      # output(32*5*5) 这里将tensor展平至一维
        x = Func.relu(self.fc1(x))      # input 32*5*5 output 120
        x = Func.relu(self.fc2(x))      # input 120 output 84
        x = self.fc3(x)                 # input 84 output 10
        return x
